package gold.contest;

import gold.utils.InputUtil;

import java.util.*;

/**
 * Created by fanzhenyu02 on 2021/12/10.
 * common problem solver template.
 */
public class LC6031 {
    public long startExecuteTime = System.currentTimeMillis();


    class Solution {
        public List<Integer> findKDistantIndices(int[] nums, int key, int k) {
            List<Integer> keyList = new ArrayList<>();
            for (int i = 0; i < nums.length; i++) {
                if (nums[i] == key) keyList.add(i);
            }

            Set<Integer> ans = new TreeSet<>();
            for (Integer pos : keyList) {
                for (int i = 0; i <= k; i++) {
                    if (pos - i >= 0 && pos - i < nums.length) ans.add(pos - i);
                    if (pos + i >= 0 && pos + i < nums.length) ans.add(pos + i);
                }
            }

            return new ArrayList<>(ans);
        }
    }

    public void run() {
        Solution solution = new Solution();
        System.out.println(solution.findKDistantIndices(InputUtil.toIntegerArray("[3,4,9,1,3,9,5]"), 9, 1));
        System.out.println(solution.findKDistantIndices(InputUtil.toIntegerArray("[2,2,2,2,2]"), 2, 2));
    }

    public static void main(String[] args) throws Exception {
        LC6031 an = new LC6031();
        an.run();

        System.out.println("\ncurrent solution total execute time: " + (System.currentTimeMillis() - an.startExecuteTime) + " ms.");
    }
}
